package com.scheduler.utils.thread;

import com.scheduler.utils.trace.Slf4jMdcUtil;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

import javax.validation.constraints.NotNull;
import java.util.Date;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.*;
import java.util.function.Supplier;

/**
 * 包装一下
 * <p>
 * 1. 可以根据提供者数据进行线程池大小控制
 */
@Slf4j
public final class MutableThreadPoolExecutor extends ThreadPoolExecutor {

    /**
     * 线程池内线程数量的提供者，方便外部修改后可以直接读取到
     */
    private final Supplier<Integer> configMaxPoolSizeSupplier;
    /**
     * 线程池名字，用来打日志的
     */
    @Getter
    private final String name;
    /**
     * 创建时间
     */
    @Getter
    private final Date createTime = new Date();

    public MutableThreadPoolExecutor(Supplier<Integer> configMaxPoolSizeSupplier, String name, int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue) {
        this(configMaxPoolSizeSupplier, name, corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, Executors.defaultThreadFactory());
    }

    public MutableThreadPoolExecutor(Supplier<Integer> configMaxPoolSizeSupplier, String name, int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory) {
        this(configMaxPoolSizeSupplier, name, corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, new AbortPolicy());
    }

    public MutableThreadPoolExecutor(Supplier<Integer> configMaxPoolSizeSupplier, String name, int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, RejectedExecutionHandler handler) {
        this(configMaxPoolSizeSupplier, name, corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, Executors.defaultThreadFactory(), handler);
    }

    /**
     * @param configMaxPoolSizeSupplier 最大线程数量的提供者
     * @param name                      for log
     */
    public MutableThreadPoolExecutor(Supplier<Integer> configMaxPoolSizeSupplier, String name, int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory, RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
        Objects.requireNonNull(name);
        Objects.requireNonNull(configMaxPoolSizeSupplier);
        this.configMaxPoolSizeSupplier = configMaxPoolSizeSupplier;
        this.name = name;
    }

    @Override
    public Future<?> submit(Runnable task) {
        return super.submit(Slf4jMdcUtil.wrapperRunnableForSubThread(task));
    }

    @Override
    public <T> Future<T> submit(Runnable task, T result) {
        return super.submit(Slf4jMdcUtil.wrapperRunnableForSubThread(task), result);
    }

    @Override
    public <T> Future<T> submit(Callable<T> task) {
        return super.submit(Slf4jMdcUtil.wrapperCallableForSubThread(task));
    }

    @Override
    public void execute(Runnable command) {
        super.execute(Slf4jMdcUtil.wrapperRunnableForSubThread(command));
    }

    /**
     * 如果提供者的限流数据修改了，则重新生成一个限流对象
     */
    @NotNull
    void renewFromConfigIfNecessary() {
        // 不用每次都执行，定时执行就好了，加个 watchDog
        int configMaxPoolSize = Optional.ofNullable(this.configMaxPoolSizeSupplier.get()).orElse(0);
        if (configMaxPoolSize <= 0) {
            return;
        }
        if (configMaxPoolSize < getCorePoolSize()) {
            log.warn("ThreadPool Change Failed, {}, {} --> {}, coreSize {}", name, getMaximumPoolSize(), configMaxPoolSize, getCorePoolSize());
            return;
        }
        if (configMaxPoolSize != getMaximumPoolSize()) {
            log.info("ThreadPool Change, {}, {} --> {}", name, getMaximumPoolSize(), configMaxPoolSize);
            try {
                // 如果核心线程数大于最大线程数，则先减少核心线程数
                if (configMaxPoolSize < getCorePoolSize()) {
                    setCorePoolSize(configMaxPoolSize);
                }
                setMaximumPoolSize(configMaxPoolSize);
            } catch (IllegalArgumentException e) {
                log.error("修改线程池大小异常，参数不正确", e);
            }
        }
    }
}
